# 参数设置
d_model = 512 # 词嵌入维度d_model
n_heads = 8
d_k = d_v = d_model //n_heads  
class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.mul_head_atten_add_norm = MultiHeadAttention_Add_Norm()
        # 多头注意力机制
        self.pos_ffn_add_norm = PoswiseFeedForwardNet_Add_Norm()
        # 提取特征
    def forward(self, enc_inputs, enc_self_attn_mask):
        '''
        enc_inputs: [batch_size, src_len, d_model]
        enc_self_attn_mask: [batch_size, src_len, src_len]
        '''
        # enc_outputs: [batch_size, src_len, d_model],
        # attn: [batch_size, n_heads, src_len, src_len] 
        # 每一个头对应一个注意力矩阵
        enc_outputs=self.mul_head_atten_add_norm(enc_inputs, 
                                                        enc_inputs, 
                                                        enc_inputs, 
                                                        enc_self_attn_mask)
        # enc_inputs to same Q,K,V
        # 乘以WQ,WK,WV生成QKV矩阵（传三个一样的参数）,
        # 因为Decoder-Encoder的mulit-head里面,
        # 需要decoder输入,encoder输出,encoder输出
        enc_outputs = self.pos_ffn_add_norm(enc_outputs)
        # enc_outputs: [batch_size, src_len, d_model]
        # 输入和输出的维度是一样的
        return enc_outputs
